import git
from pathlib import Path
import json
from typing import List, Dict
import logging
import re

class GitCodeCollector:
	def __init__(self, base_dir: str):
		self.base_dir = Path(base_dir)
		self.setup_logging()
		
	def setup_logging(self):
		logging.basicConfig(
			level=logging.INFO,
			format='%(asctime)s - %(levelname)s - %(message)s'
		)
		self.logger = logging.getLogger(__name__)
	
	def find_git_repos(self) -> List[Path]:
		"""查找所有git仓库目录"""
		git_repos = []
		for item in self.base_dir.iterdir():
			if item.is_dir() and (item / '.git').exists():
				git_repos.append(item)
		return git_repos

	def list_files_and_dirs(self, repo_path: Path, path: Path = None) -> Dict:
		"""递归列出目录下的所有文件和子目录，构建树状结构"""
		if path is None:
			path = repo_path
			
		# 创建当前目录节点
		current_node = {
			'name': path.name if path != repo_path else repo_path.name,
			'path': str(path.relative_to(repo_path)) if path != repo_path else '',
			'type': 'dir',
			'children': []
		}
		
		try:
			# 列出当前目录下的所有项目
			for item in path.iterdir():
				if item.is_dir():
					if item.name != '.git':  # 排除 .git 目录
						print(f"添加目录 {item.relative_to(repo_path)}")
						# 递归处理子目录
						sub_tree = self.list_files_and_dirs(repo_path, item)
						current_node['children'].append(sub_tree)
				else:
					print(f"添加文件 {item.relative_to(repo_path)}")
					# 添加文件节点
					file_node = {
						'name': item.name,
						'path': str(item.relative_to(repo_path)),
						'type': 'file'
					}
					current_node['children'].append(file_node)
					
		except PermissionError:
			self.logger.warning(f"无法访问目录: {path}")
			return current_node

		return current_node

	def get_file_history(self, repo: git.Repo, file_path: Path) -> List[Dict]:
		"""获取文件的 Git 提交历史"""
		try:
			# 获取文件的提交历史
			commits = list(repo.iter_commits(paths=str(file_path.relative_to(repo.working_dir))))
			history = []
			
			for commit in commits:
				history.append({
					'hash': commit.hexsha,
					'author': commit.author.name,
					'date': commit.authored_datetime.isoformat(),
					'message': commit.message.strip()
				})
			
			return history
		except Exception as e:
			self.logger.error(f"获取文件历史时出错 {file_path}: {e}")
			return []

	def collect_code_files(self, repo_path: Path) -> List[Dict]:
		"""收集单个仓库中的代码文件"""
		code_files = []
		
		try:
			repo = git.Repo(repo_path)
		except git.exc.InvalidGitRepositoryError:
			self.logger.warning(f"无效的 Git 仓库: {repo_path}")
			return []
		
		# 从repo用git对象获取仓库远程全名
		remote_url = repo.remote().url
		
		# 获取目录树结构
		file_tree = self.list_files_and_dirs(repo_path)
		
		# 支持的文件类型
		supported_extensions = {'.py', '.cpp', '.c', '.h', '.java', '.js', 
							  '.go', '.xml', '.yml', '.md', '.sh', '.proto'}
		
		# 递归处理目录树，添加文件内容和历史
		def process_node(node):
			if node['type'] == 'file':
				file_path = repo_path / node['path']
				if file_path.suffix in supported_extensions:
					try:
						with open(file_path, 'r', encoding='utf-8') as f:
							content = f.read()
						
						# 获取文件的 Git 历史
						history = self.get_file_history(repo, file_path)
						
						code_files.append({
							'repo': remote_url,
							'path': "/" + node['path'],
							'name': node['name'],
							'content': content,
							'suffix': file_path.suffix,
							'type': 'file',
							'history': [x['message'] for x in history[:5]]
						})
					except UnicodeDecodeError:
						self.logger.warning(f"无法解码文件: {file_path}")
					except Exception as e:
						self.logger.error(f"处理文件时出错 {file_path}: {e}")
			if node['type'] == 'dir':  # dir type
				code_files.append({
					'repo': remote_url,
					'path': "/" + node['path'],
					'name': node['name'],
					'type': 'dir',
					'children': node['children']  # 保持完整的子节点信息
				})
				# 递归处理子节点
				for child in node['children']:
					process_node(child)
		
		process_node(file_tree)
		return code_files

	def save_dataset(self, code_files: List[Dict], output_file: str):
		"""保存收集的代码到数据集文件"""
		dataset = []
		
		for code_file in code_files:
			if code_file["type"] == "file":
				# 尝试从代码中提取注释作为功能描述
				description = self.extract_description_from_code(code_file["content"])
				
				# 构建训练样本格式
				sample = {
					'conversations': [
						{
							'role': 'system',
							'content': '你是NCT公司Engineering-Server部门的智能代码助手，了解Engineering-Server的代码结构，知道每个代码文件的内容。'
						},
						{
							'role': 'human', 
							'content': code_file["repo"] + '工程的文件:```' + code_file["path"] + '```的内容是什么？'
						},
						{
							'role': 'assistant',
							'content': '这个文件的类型是:```' + code_file["suffix"][1:] + '```\n' + '内容是:\n```\n' + code_file["content"] + '\n```\n' + 
								'文件git历史:\n```json\n' + json.dumps(code_file["history"], ensure_ascii=False, indent=2) + '\n```\n'
						}
					]
				}
				dataset.append(sample)
			elif code_file["type"] == "dir":
				# 从完整的子节点信息中分离目录和文件
				dirs = []
				files = []
				for child in code_file["children"]:
					if child["type"] == "dir":
						dirs.append(child["name"])
					else:
						files.append(child["name"])
				
				sample = {
					'conversations': [
						{
							'role': 'system',
							'content': '你是NCT公司Engineering-Server部门的智能代码助手，了解Engineering-Server的代码结构，知道每个代码文件的内容。'
						},
						{
							'role': 'human',
							'content': code_file["repo"] + '工程的目录:```' + code_file["path"] + '```包含什么？'
						},
						{
							'role': 'assistant',
							'content': f'目录 `{code_file["name"]}` 包含以下目录:```\n' + '\n'.join([f'- {d}' for d in dirs]) + '\n```\n' + 
								'包含以下文件:```\n' + '\n'.join([f'- {f}' for f in files]) + '\n```\n'
						}
					]
				}
				dataset.append(sample)
		
		with open(output_file, 'w', encoding='utf-8') as f:
			json.dump(dataset, f, ensure_ascii=False, indent=2)
		
		self.logger.info(f'数据集已保存到: {output_file}')

	def extract_description_from_code(self, code: str) -> str:
		"""从代码中提取注释作为功能描述"""
		# 提取Python和Java的注释
		python_comments = re.findall(r'#.*', code)  # Python单行注释
		java_single_comments = re.findall(r'//.*', code)  # Java单行注释
		java_multi_comments = re.findall(r'/\*.*?\*/', code, flags=re.DOTALL)  # Java多行注释
		
		# 合并所有注释
		all_comments = python_comments + java_single_comments + java_multi_comments
		# 清理注释内容
		cleaned_comments = [comment.strip() for comment in all_comments]
		return ' '.join(cleaned_comments) if cleaned_comments else '未提供详细描述'

	def run(self, output_file: str):
		"""运行代码收集流程"""
		self.logger.info(f'开始扫描目录: {self.base_dir}')
		
		all_code_files = []
		repos = self.find_git_repos()
		
		for repo in repos:
			self.logger.info(f'处理仓库: {repo.name}')
			code_files = self.collect_code_files(repo)
			all_code_files.extend(code_files)
			
		self.logger.info(f'共收集了 {len(all_code_files)} 个代码文件')
		self.save_dataset(all_code_files, output_file) 
		self.logger.info(f'数据集已保存到: {output_file}')